import unittest

import numpy as np
import torch

from training.tools.metrics import _fast_hist, overall_pixel_accuracy, iinc

shape = (1, 2, 2)


class MetricsTestCase(unittest.TestCase):

    def test_metrics(self):
        target = np.array([
            [1, 1, 1, 1, 1],
            [1, 1, 1, 1, 1],
            [1, 1, 1, 1, 1],
            [1, 1, 1, 1, 1],
            [1, 1, 1, 1, 1],
        ])
        target = torch.from_numpy(target)

        predict = np.array([
            [1, 1, 1, 1, 1],
            [1, 0, 0, 0, 1],
            [1, 0, 0, 0, 1],
            [1, 0, 0, 0, 1],
            [1, 1, 1, 1, 1],
        ])
        predict = torch.from_numpy(predict)

        print(target.shape, predict.shape)
        hist = _fast_hist(target, predict, 2)
        overall_acc = overall_pixel_accuracy(hist)
        print(overall_acc)
        self.assertEqual(0.64 * 100, overall_acc)

    def test_iinc_a(self):
        pred = torch.zeros(shape)
        gt = torch.zeros(shape)
        self.assertEqual(0.00, iinc(gt, pred))

    def test_iinc_b(self):
        pred = torch.zeros(shape)
        pred[0, 0, 0] = 1
        gt = torch.zeros(shape)
        self.assertEqual(1 / 2.75, iinc(gt, pred))

    def test_iinc_c(self):
        pred = torch.zeros(shape)
        pred[0, 0] = 1
        gt = torch.zeros(shape)
        self.assertEqual(1 / 2.5, iinc(gt, pred))

    def test_iinc_d(self):
        pred = torch.zeros(shape)
        pred[0, 0] = 1
        gt = torch.zeros(shape)
        gt[0, 1, 0] = 1
        self.assertEqual(2 / 2.25, iinc(gt, pred))

    def test_iinc_e(self):
        pred = torch.zeros(shape)
        pred[0, 0] = 1
        pred[0, :, 0] = 1
        gt = torch.zeros(shape)
        gt[0, 1, 0] = 1
        self.assertEqual((1 / 2.25 * 2 / 3), iinc(gt, pred))

    def test_iinc_f(self):
        pred = torch.zeros(shape)
        pred[0, 0] = 1
        pred[0, :, 0] = 1
        gt = torch.zeros(shape)
        gt[0, 1] = 1
        gt[0, :, 1] = 1
        self.assertEqual((1 / 2) * (2 / 3), iinc(gt, pred))

    def test_iinc_g(self):
        pred = torch.ones(shape)
        gt = torch.zeros(shape)
        gt[0, :, 0] = 1
        self.assertEqual((1 / 2) * (1 / 2), iinc(gt, pred))


if __name__ == '__main__':
    unittest.main()
